from fastapi import (
    APIRouter, 
    Query, 
    Header,
    Cookie, 
    Path, 
    Body, 
    File,
    Form,
    status,
    Depends,
    UploadFile,
    BackgroundTasks,
    HTTPException,
    Response,
)
from fastapi.responses import (
    FileResponse,
    HTMLResponse,
    JSONResponse,
    ORJSONResponse,
    PlainTextResponse,
    RedirectResponse,
    Response,
    StreamingResponse,
    UJSONResponse,
)
from fastapi.openapi.docs import (
    get_redoc_html,
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
)
from pydantic import BaseModel, Field #, Required
from typing import Union,Any,List,Optional
from starlette.background import BackgroundTask
from functools import lru_cache
from mimetypes import guess_type
import urllib.parse
from osgeo import gdal,ogr
import os

import toolbox


router = APIRouter(
    prefix="/osgeotoolbox",
    tags=["osgeotoolbox"],
    dependencies=[ 
        # Depends(get_token_header),
        # Depends(get_query_token),
    ],
    responses={404: {"description": "Not found"}},
)


def file_iterator(file_path, offset, chunk_size):
    """
    文件生成器
    :param file_path: 文件绝对路径
    :param offset: 文件读取的起始位置
    :param chunk_size: 文件读取的块大小
    :return: yield
    """
    with open(file_path, 'rb') as f:
        f.seek(offset, os.SEEK_SET)
        while True:
            data = f.read(chunk_size)
            if data:
                yield data
            else:
                break


# 统计村域作物种植面积
@lru_cache(maxsize=300)
@router.post(
    "/uploadfile_statzwarea",
    summary="根据上传的文件进行，统计村域作物种植面积",
    description="统计村域作物种植面积",  # 与函数注释互斥
    response_description="统计村域作物种植面积",
    # response_model=FileResponse,
    deprecated=False,  # 是否启用
)
async def stat_zw_area(
    zw_year: str = Form(
        default="2024",
        title="作物种植年份",
        description="作物种植年份必须指定",
        deprecated=False,  # 是否启用
    ), 
    zwdm: str = Form(
        default="07",
        title="作物代码",
        description="作物代码必须指定",
        deprecated=False,  # 是否启用
    ),
    zwmc: str = Form(
        default="冬小麦",
        title="作物名称",
        description="作物名称必须指定",
        deprecated=False,  # 是否启用
    ),
    jcsj: str = Form(
        default="2023-04-01",
        title="年月日",
        description="年月日必须指定",
        deprecated=False,  # 是否启用
    ),
    rspdid: str = Form(
        default="zzjg_411502_07_2024_20240320",
        title="任务编码",
        description="任务编码必须指定",
        deprecated=False,  # 是否启用
    ),
    qh_shp_files: List[UploadFile] = File(
        title="区划矢量shp文件",
        description="矢量shp文件",
        media_type="multipart/form-data",
    ),
    zw_shp_files: List[UploadFile] = File(
        title="作物矢量shp文件",
        description="矢量shp文件",
        media_type="multipart/form-data",
    ),
):  
    save_dir = r"/data/fengyy/dateset/temp"
    for shp_file in qh_shp_files:
   
        tmp_path = os.path.join(save_dir,shp_file.filename)
        if os.path.splitext(shp_file.filename)[-1]==".shp":
            qh_shp_path=tmp_path
        if not os.path.exists(tmp_path):
            with open(tmp_path,"wb") as file_handle:
                for chunk in iter(lambda: shp_file.file.read(1024), b''):
                    file_handle.write(chunk)
        await shp_file.close()
    for shp_file in zw_shp_files:

        tmp_path = os.path.join(save_dir,shp_file.filename)
        if os.path.splitext(shp_file.filename)[-1]==".shp":
            zw_shp_file=tmp_path
        if not os.path.exists(tmp_path):
            with open(tmp_path,"wb") as file_handle:
                for chunk in iter(lambda: shp_file.file.read(1024), b''):
                    file_handle.write(chunk)
        await shp_file.close()
    res_shp_path = os.path.join(save_dir,"result.shp")
    dbf_path = toolbox.zwcy_2_tb(
        zw_shp_file,
        qh_shp_path,
        res_shp_path,
        zw_year,
        zwdm,
        zwmc,
        jcsj,
        rspdid
    )
    filename = "种植结构统计表.dbf"
    encoded_filename = urllib.parse.quote(filename)
    headers = {'content-disposition': f'attachment; filename="{encoded_filename}"'}
    return FileResponse(
        dbf_path,
        # filename=encoded_filename,
        status_code=200,
        media_type='gzip',
        background=BackgroundTask(lambda: os.remove(dbf_path)),
        headers=headers,
        stat_result=None,
        method=None,
        content_disposition_type="attachment",
    )


# 根据上传的栅格裁剪矢量
@lru_cache(maxsize=300)
@router.post(
    "/uploadfile_shpcliptif",
    summary="根据上传的文件进行，矢量裁剪栅格",
    description="需要矢量文件 栅格文件 保存目录",  # 与函数注释互斥
    response_description="矢量裁剪栅格的结果",
    # response_model=FileResponse,
    deprecated=False,  # 是否启用
)
async def upload_shp_clip_tif(
    shp_files: List[UploadFile] = File(
        title="矢量shp文件",
        description="矢量shp文件",
        media_type="multipart/form-data",
    ),
    tif_file: UploadFile = File(
        title="栅格tif文件",
        description="栅格tif文件",
        media_type="multipart/form-data",
    ),
    save_dir: str = Form(default=r'/home/fengyang/data/202401/202403/upload'),
) -> FileResponse:
    shp_path = ""
    for shp_file in shp_files:
        tmp_path = os.path.join(save_dir,shp_file.filename)
        if os.path.splitext(shp_file.filename)[-1]==".shp":
            shp_path=tmp_path
        if not os.path.exists(tmp_path):
            with open(tmp_path,"wb") as file_handle:
                for chunk in iter(lambda: shp_file.file.read(1024), b''):
                    file_handle.write(chunk)
        await shp_file.close()
    tif_path = os.path.join(save_dir,tif_file.filename)
    if not os.path.exists(tif_path):
        with open(tif_path,"wb") as file_handle:
            for chunk in iter(lambda: tif_file.file.read(1024), b''):
                file_handle.write(chunk)
    await tif_file.close()
    dst_path = os.path.join(save_dir,"result.tif")
    dataset = ogr.Open(shp_path)
    layer = dataset.GetLayer()
    layer_name = layer.GetName()
    dataset.Destroy()
    gdal.Warp(
        dst_path, 
        tif_path, 
        format="GTiff", # JPEG GTiff
        cutlineDSName=shp_path,
        cutlineLayer=layer_name,
        cutlineWhere=None, 
        cutlineBlend=True,
        cropToCutline=True,
        resampleAlg=gdal.GRIORA_NearestNeighbour, # GRIORA_NearestNeighbour,GRIORA_Bilinear
        dstNodata=None,
        warpMemoryLimit=1024*1024*1024*50,
        multithread=10)
    return FileResponse(
        dst_path,
        status_code=200,
        headers=None,
        media_type=None,
        background=None,
        filename=dst_path,
        stat_result=None,
        method=None,
        content_disposition_type="attachment",
    )